import os
from tree_sitter_c_config import rust_parser, c_parser

# 提取extern C声明部分的代码
# 传入文件路径，返回 str
def extract_extern_declaration(path:str):
    if not os.path.exists(path):
        return None
    with open (path) as file:
        code = file.read()
    if code.strip() == '':
        return None
    
    parser = rust_parser
    tree = parser.parse(bytes(code, 'utf8'))
    # 查找extern C声明部分的代码
    def traverse(node):
        if node.type == 'foreign_mod_item':
            start_byte = node.start_byte
            end_byte = node.end_byte
            return code[start_byte:end_byte]
        for child in node.children:
            res = traverse(child)
            if res:
                return res
    result = traverse(tree.root_node) 
    return result if result else ""

# 从extern C的声明块中逐个提取函数签名
# 传入code，返回dict（函数名到签名代码的映射）
def extract_func_signature_from_extern_declaration(code:str):
    if code.strip() == '':
        return None
    
    # 加载parser
    parser = rust_parser
    # 解析代码
    tree = parser.parse(bytes(code, 'utf8'))
    results = {}
    def traverse(node):
        if node.type == 'function_signature_item':
            start_byte = node.start_byte
            end_byte = node.end_byte
            sig_code = code[start_byte:end_byte]

            # 这文件把tree-sitter干error了，绝了
            if "te_interp" in sig_code:
                pos = sig_code.find("fn")
                sig_code = sig_code[pos:]

            # 获取函数名
            for child in node.children:
                if child.type == 'identifier':
                    function_name = code[child.start_byte:child.end_byte]
                    results[function_name] = sig_code
                    break
        for child in node.children:
            traverse(child)

    traverse(tree.root_node) 
    return results
        
# 提取全局结构体定义
# 传入文件路径，返回 str[]
def extract_struct_declaration(path:str):
    if not os.path.exists(path):
        return None
    # 解析代码
    with open (path) as file:
        code = file.read()
    if code.strip() == '':
        return None
    
    # 加载parser
    parser = rust_parser
    tree = parser.parse(bytes(code, 'utf8'))
    results = []

    def traverse(node):
        if node.type == 'struct_item':
            start_byte = node.start_byte
            end_byte = node.end_byte
            results.append(code[start_byte:end_byte])
        for child in node.children:
            traverse(child)
    traverse(tree.root_node) 
    return results

# 提取静态变量定义
# 传入文件路径，返回 str[]
def extract_static_declaration(path:str):
    if not os.path.exists(path):
        return None
    # 解析代码
    with open (path) as file:
        code = file.read()
    if code.strip() == '':
        return None
    
    # 加载parser
    parser = rust_parser
    tree = parser.parse(bytes(code, 'utf8'))
    results = []
    def traverse(node):
        if node.type == 'static_item':
            start_byte = node.start_byte
            end_byte = node.end_byte
            # 过滤tinyexpr的特例
            if "static mut stdout" not in code[start_byte:end_byte]:
                results.append(code[start_byte:end_byte])
        for child in node.children:
            traverse(child)
    traverse(tree.root_node) 
    return results

# 提取静态类型定义
# 传入文件路径，返回 str[]
def extract_type_declaration(path:str):
    if not os.path.exists(path):
        return None
    with open (path) as file:
        code = file.read()
    if code.strip() == '':
        return None
    
    # 加载parser
    parser = rust_parser
    tree = parser.parse(bytes(code, 'utf8'))
    results = []

    def traverse(node):
        if node.type == 'type_item':
            start_byte = node.start_byte
            end_byte = node.end_byte
            results.append(code[start_byte:end_byte])
        for child in node.children:
            traverse(child)
    traverse(tree.root_node) 
    return results

# 提取所有函数名
# 传入文件路径，返回 str[]
def extract_function_names(path:str):
    if not os.path.exists(path):
        return None
    with open (path) as file:
        code = file.read()
    if code.strip() == '':
        return []
    
    # 加载parser
    parser = rust_parser
    tree = parser.parse(bytes(code, 'utf8'))
    results = []
    def traverse(node):
        if node.type == 'function_item':
            # 获取函数名
            for child in node.children:
                if child.type == 'identifier':
                    function_name = code[child.start_byte:child.end_byte]
                    results.append(function_name)
                    break
        for child in node.children:
            traverse(child)
    traverse(tree.root_node) 
    return results

# 提取所有c label的块
# 传入文件路径，返回 str[]
def extract_c_label_blocks(path:str):
    if not os.path.exists(path):
        return None
    with open (path) as file:
        code = file.read()
    if code.strip() == '':
        return []
    
    # 加载parser
    parser = rust_parser
    tree = parser.parse(bytes(code, 'utf8'))
    results = []
    def traverse(node):
        if node.type == 'expression_statement':
            # print('-' * 40)
            
            for child in node.children:
                if child.type == 'block':
                    block_node = child
                    for c in block_node.children:
                        if c.type == "label":
                            label_node = c
                            if "'c_" in label_node.text.decode('utf-8'):
                                results.append(code[node.start_byte:node.end_byte])
                    
        for child in node.children:
            traverse(child)   

    traverse(tree.root_node) 
    return results      

def extract_call_func_names(code:str):
    if code.strip() == '':
        return None
    
    # 加载parser
    parser = rust_parser
    # 解析代码
    tree = parser.parse(bytes(code, 'utf8'))
    results = []
    def traverse(node):
        if node.type == 'call_expression':
            func_node = node.child_by_field_name("function")
            start_byte = func_node.start_byte
            end_byte = func_node.end_byte
            func_name = code[start_byte:end_byte]
            results.append(func_name)

        for child in node.children:
            traverse(child)

    traverse(tree.root_node) 
    return results

def extract_assert_macro(path:str):
    if not os.path.exists(path):
        return None
    with open (path) as file:
        code = file.read()
    if code.strip() == '':
        return ""
    
    parser = c_parser
    tree = parser.parse(bytes(code, 'utf8'))
    results = ""
    def traverse(node):
        nonlocal results
        if node.type == 'preproc_function_def':
            child = node.child_by_field_name('name')
            if code[child.start_byte:child.end_byte] == 'assert':
                if ".cpp" in path:
                    results = code[node.start_byte:node.end_byte]
                else:
                    index = node.parent.children.index(node)
                    node3 = node.parent.children[index+2]
                    results = code[node.start_byte:node3.end_byte]

        for child in node.children:
            traverse(child)
    traverse(tree.root_node) 
    return results
if __name__ == "__main__":
    # path = "/mnt/sda/xc/C2Rust/c_rust_agents/vivo-c2rust/output/primary/tests/test_arraylist.rs"
    # with open (path) as file:
    #     code = file.read()

    # extern_code = extract_extern_declaration(path)
    # print(extern_code)
    # results = extract_func_signature_from_extern_declaration(extern_code)
    # for key in results.keys():
    #     print(key)
    #     print(results[key])
    #     if results[key] in code:
    #         print("find")

    # results = extract_struct_declaration(path)
    # for res in results:
    #     print(res)
    #     if res in code:
    #         print("find")

    # results = extract_static_declaration(path)
    # # print(results)
    # for res in results:
    #     print(res)
    #     if res in code:
    #         print("find")



    # results = extract_type_declaration(path)
    # # print(results)
    # for res in results:
    #     print(res)
    #     if res in code:
    #         print("find")
    

    # path = "/mnt/sda/xc/C2Rust/c_rust_agents/vivo-c2rust/cal-safe-ratio/primary/src/arraylist.rs"
    # with open (path) as file:
    #     code = file.read()
    # results = extract_function_names(code)
    # for res in results:
    #     print(res)
    #     if res in code:
    #         print("find")
    

    # file_name = "test_avl_tree.rs"
    # start_index = len("test_")
    # end_index = file_name.find(".rs")
    # module_name = file_name[start_index:end_index]
    # print(module_name)

    # results = extract_extern_declaration('./test.rs')
    # print(results)


    code ="""
pub extern "C" fn test_avl_tree_insert_lookup() {
    unsafe {
        alloc_test_set_limit(-1);
        let mut tree: *mut AVLTree = std::ptr::null_mut();
        let mut node: *mut AVLTreeNode = std::ptr::null_mut();
        let mut i: libc::c_uint = 0;
        let mut value: *mut libc::c_int = std::ptr::null_mut();

        tree = avl_tree_new(::core::mem::transmute::<
            Option<unsafe extern "C" fn(*mut libc::c_void, *mut libc::c_void) -> libc::c_int>,
            AVLTreeCompareFunc,
        >(Some(
            int_compare
                as unsafe extern "C" fn(*mut libc::c_void, *mut libc::c_void) -> libc::c_int,
        )));

        for i in 0..1000 {
            test_array[i as usize] = i as libc::c_int;
            avl_tree_insert(
                tree,
                &mut test_array[i as usize] as *mut libc::c_int as AVLTreeKey,
                &mut test_array[i as usize] as *mut libc::c_int as AVLTreeValue,
            );
            num_assert = num_assert.wrapping_add(1);

            assert!(
                avl_tree_num_entries(tree) == i + 1,
                "avl_tree_num_entries(tree) == i + 1"
            );
            validate_tree(tree);
        }

        assert!(
            !avl_tree_root_node(tree).is_null(),
            "avl_tree_root_node(tree) != NULL"
        );

        for i in 0..1000 {
            let mut key = i as libc::c_uint;
            node = avl_tree_lookup_node(tree, &mut key as *mut libc::c_uint as AVLTreeKey);
            assert!(!node.is_null(), "node != NULL");

            value = avl_tree_node_key(node) as *mut libc::c_int;
            assert!(*value == i as libc::c_int, "*value == (int) i");

            value = avl_tree_node_value(node) as *mut libc::c_int;
            assert!(*value == i as libc::c_int, "*value == (int) i");
        }

        let mut key = 1100 as libc::c_uint;
        assert!(
            avl_tree_lookup_node(tree, &mut key as *mut libc::c_uint as AVLTreeKey).is_null(),
            "avl_tree_lookup_node(tree, &i) == NULL"
        );

        avl_tree_free(tree);
    }
}
"""
    res = extract_call_func_names(code)
    print(res)